Analysis of Workflow of MainDiffusionV1 in verl-omni
Analysis of the main_diffusion_v1 entrance of v1 trainer.
Table of Contents
The bash examples/......sh file in examples are actually a wrapper over verl_omni.trainer.main_diffusion.py or main_diffusion_v1.py.
1. The Main Entrance
The entrance is at main_diffusion_v1.py:L136. It first sets up devices, resolves the configuration using OmegaConf, validates configurations, and validates attention consistency.
Then, based on the use_v1 switch in configuration, it routes to run_diffusion_v1() or legacy run_diffusion() function. In this post, we follow the run_diffusion_v1() path.
2. run_diffusion_v1()
Although this function takes a config parameter and a task_runner_class parameter, the second parameter is always set to None when invoked within main(), so the task_runner_class is always initialized to the DiffusionTaskRunnerV1 class1. After initializing the task_runner_class class, we initialize the runner, which is a handle in Ray to the object managed by Ray and is obtained through remote() method.
Footnotes:
main_diffusion_v1.py:L87